from APP.config.log_conf import log
import httpx
from APP.config.conf import settings
from typing import AsyncGenerator
import json
import asyncio
import re
import pandas as pd
import plotly.figure_factory as ff
from APP.common.middleware.vanna_api.userVanna import userVanna

class ChatDataServices:
    
    @staticmethod
    #请求dbgpt接口，获取响应结果
    async def request_dbgpt(message: str,model:str,chat_param:str,temperature:float) -> dict:
        #请求头
        headers = {
            "accept": "application/json",
            "Content-Type": "application/json",
            "Authorization": "Bearer dbgpt"
        }

        #请求体
        payload = {
            "messages": message,
            "model": model,
            "chat_mode": "chat_data",
            "chat_param": chat_param,
            "temperature": temperature
        }

        async with httpx.AsyncClient() as client:
            response = await client.post(
                url=settings.DBGPT_API_URL,
                headers=headers,
                json=payload,
                timeout=None
            )

            response_data = response.json()
            # 打印响应数据
            log.info(f"Response data: {response_data}")

            return response_data
        
    # 处理响应结果，提取数据
    @staticmethod
    async def generate_stream_response(message:str,response_data: dict) -> AsyncGenerator[str, None]:
        res = response_data['choices'][0]['message']
        #回复话内容
        content = res['content'].split('\\n')[0]

        for char in content:
            yield f"data: {json.dumps({'message': char})}\n\n"
            await asyncio.sleep(0.05)  # 控制输出速度
        
        #获取chart-view content="开头，"/>结尾的字符串
        pattern = r'chart-view content="(.*?)" />'
        match = re.search(pattern, res['content'].split('\\n')[-1])
        
        if match:
            # 获取匹配到的字符串
            json_str = match.group(1)
            #将字符串转换成字典
            data_dict = json.loads(json_str)
            sql = data_dict['sql']
            #获取查询出来的数据
            df_str = data_dict['data']
            #将列表转成dataframe
            df = pd.DataFrame(df_str)
            #生成表格
            fig_table =ff.create_table(df)

            # await asyncio.sleep(1)  # 控制输出速度

            #vanna自动生成图表
            instance = userVanna()
            fig_echarts = await instance.ask(message,df,sql)
              
            # 发送最终数据
            final_data = {
                "table": fig_table.to_json(),
                "sql": sql,
                "echarts": fig_echarts  
            }
            log.info(f"chatdata_final_data: {final_data}")
            yield f"data: {json.dumps(final_data)}\n\n"
        else:
            yield f"data: {json.dumps({'interrupt': '没有生成图表'})}\n\n"


ChatDataService: ChatDataServices = ChatDataServices()